package main

import (
	"fmt"
	"log"
	"sort"
	"strconv"
	"strings"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
)

var (
	commandMap = make(map[string]int)    
	ackMap     = make(map[string]string) 
	ackLenMap  = make(map[string]int)   

)

type Pair struct {
	Key   string
	Value uint64
}

func formatBytes(bytes uint64) string {
	const unit = 1024
	if bytes < unit {
		return fmt.Sprintf("%d B", bytes)
	}
	div, exp := unit, 0
	for n := bytes / unit; n >= unit; n /= unit {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.1f %cB", float64(bytes)/float64(div), "KMGTPE"[exp])
}

func printPacketInfo(packet gopacket.Packet) {
	quiteLayer := []string{"ethernetLayer", "All"} //"ethernetLayer", "ipLayer", "tcpLayer","applicationLayer", "All"
	qfunc := func(searchStr string) bool {
		found := false
		for _, str := range quiteLayer {
			if str == searchStr {
				found = true
				break
			}
		}
		return found
	}
	if !qfunc("All") {
		// Iterate over all layers, printing out each layer type
		fmt.Println("All packet layers:")
		for _, layer := range packet.Layers() {
			fmt.Println("- ", layer.LayerType()) //通常来说都是Ethernet、IPv4、TCP、Payload四个
		}
	}
	// 判断layer是否存在错误
	if err := packet.ErrorLayer(); err != nil {
		fmt.Println("Error decoding some part of the packet:", err)
	}
	// 判断数据包是否为以太网数据包，可解析出源mac地址、目的mac地址、以太网类型（如ip类型）等
	ethernetLayer := packet.Layer(layers.LayerTypeEthernet)
	if ethernetLayer != nil && !qfunc("ethernetLayer") {
		fmt.Println("Ethernet layer detected.")
		ethernetPacket, _ := ethernetLayer.(*layers.Ethernet)
		fmt.Println("Source MAC: ", ethernetPacket.SrcMAC)
		fmt.Println("Destination MAC: ", ethernetPacket.DstMAC)
		fmt.Println("Ethernet type: ", ethernetPacket.EthernetType)
	}
	// 判断数据包是否为IP数据包，可解析出源ip、目的ip、协议号等
	ipLayer := packet.Layer(layers.LayerTypeIPv4)
	var ip *layers.IPv4
	if ipLayer != nil && !qfunc("ipLayer") {
		ip, _ = ipLayer.(*layers.IPv4)
		// TOS, Length, Id, Flags, FragOffset, TTL, Protocol (TCP?),
		// Checksum, SrcIP, DstIP
		fmt.Printf("From %s to %s\t Protocol %s \n", ip.SrcIP, ip.DstIP, ip.Protocol)
	}
	// 判断数据包是否为TCP数据包，可解析源端口、目的端口、seq序列号、tcp标志位等
	tcpLayer := packet.Layer(layers.LayerTypeTCP)
	var tcp *layers.TCP
	if tcpLayer != nil && !qfunc("ipLayer") {
		// fmt.Println("TCP layer detected.")
		tcp, _ = tcpLayer.(*layers.TCP)
		// TCP layer variables:
		// SrcPort, DstPort, Seq, Ack, DataOffset, Window, Checksum, Urgent
		// Bool flags: FIN, SYN, RST, PSH, ACK, URG, ECE, CWR, NS
		fmt.Printf("From port %d to %d\t,Seq Ack: %d %d \n", tcp.SrcPort, tcp.DstPort, tcp.Seq, tcp.Ack)
	}
	// 从数据包的各个层次中提取数据
	applicationLayer := packet.ApplicationLayer()
	if applicationLayer != nil && !qfunc("applicationLayer") {
		// 获取数据的字节数组
		data := applicationLayer.Payload()
		// 将字节数组转换为字符串或根据需要进行处理
		dataStr := string(data)
		substr := "\r\n"
		if tcp.PSH && tcp.DstPort == 6379 {
			var resStr []string
			dataStrList := strings.Split(dataStr, substr)[1:]
			for _, str := range dataStrList {
				if !strings.HasPrefix(str, "$") && str != "" {
					resStr = append(resStr, str)
				}
			}
			finnalCommand := strings.Join(resStr, " ")
			commandMap[finnalCommand]++
			commandMap["Total"]++
			ackMap[strconv.Itoa(int(tcp.Seq))] = finnalCommand
		} else if tcp.SrcPort == 6379 {
			ackLenMap[strconv.Itoa(int(tcp.Ack))] += len(data)
			ackLenMap["Total"] += len(data)
		}

	}
}

func SortkvMap(m map[string]int) []Pair {
	pairs := make([]Pair, 0, len(m))
	for key, value := range m {
		pairs = append(pairs, Pair{key, uint64(value)})
	}
	sort.Slice(pairs, func(i, j int) bool {
		return pairs[i].Value < pairs[j].Value
	})
	return pairs
}
func main() {
	capFile := "C:\\Users\\PC\\Desktop\\redis.cap" //读包位置
	handle, err := pcap.OpenOffline(capFile)
	if err != nil {
		log.Fatal(err)
	}
	defer handle.Close()
	// 设置过滤条件
	// filter := "tcp src port 42176 and tcp[13:1] > 0"
	// err = handle.SetBPFFilter(filter)
	// if err != nil {
	// 	log.Fatal(err)
	// }
	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
	for packet := range packetSource.Packets() {
		printPacketInfo(packet)
	}
	// sortedComand := SortkvMap(commandMap)
	// for _, pair := range sortedComand {
	// 	fmt.Printf("%v\n", pair)
	// }
	sortedAck := SortkvMap(ackLenMap)
	for _, paira := range sortedAck {
		fmt.Printf("ack_row: %s,大小: %s,命令： %s ,命令执行次数: %d  \n", paira.Key, formatBytes(paira.Value), ackMap[paira.Key], commandMap[ackMap[paira.Key]])
	}
}
